import cv2
from djitellopy import tello
import KeyPressModule as kp
import time
import numpy as np
import math

######## PARAETERS ########
fSpeed = 117 / 10
aSpeed = 360 / 10
interval = 0.25

dInterval = fSpeed * interval
aInterval = aSpeed * interval
###########################
x, y = 500, 500
a = 0
yaw = 0
points = [(0,0), (0,0)]

kp.init()
me = tello.Tello()
me.connect()
print(me.get_battery())

me.streamon()
global img

def getKeyBoardInput():
    lr, fb, ud, yv = 0, 0, 0, 0
    global yaw, x, y, a
    speed = 15
    aspeed = 50
    d = 0


    if kp.getKey("LEFT"):
        lr = -speed
        d = -dInterval
        a = -180

    elif kp.getKey("RIGHT"):
        lr = speed
        d = dInterval
        a = 180


    if kp.getKey("UP"):
        fb = speed
        d = dInterval
        a = 270

    elif kp.getKey("DOWN"):
        fb = -speed
        d = -dInterval
        a = -90


    if kp.getKey("w"):
        ud = speed
    elif kp.getKey("s"):
        ud = -speed

    if kp.getKey("a"):
        yv = -aspeed
        yaw -= aInterval

    elif kp.getKey("d"):
        yv = aspeed
        yaw += aInterval

    if kp.getKey("q"):
        me.takeoff()
        time.sleep(3)

    if kp.getKey("e"):
        me.land()

    if kp.getKey("z"):
        cv2.imwrite(f'Resources/Images/{time.time()}.jpg', img)
        time.sleep(0.3)

    time.sleep(interval)
    a += yaw
    x += int(d*math.cos(math.radians(a)))
    y += int(d*math.sin(math.radians(a)))

    return [lr, fb, ud, yv, x, y]


def detect_circle_by_fitting(frame):
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # ⭐ 1. 直方图均衡化（增强对比度）
    gray = cv2.equalizeHist(gray)

    # ⭐ 2. 轻度去噪（避免模糊圆）
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)

    # ⭐ 3. Canny 边缘检测
    edges = cv2.Canny(blurred, 50, 150)

    # ⭐ 4. 轮廓检测
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    for cnt in contours:
        # 计算拟合圆
        (x, y), radius = cv2.minEnclosingCircle(cnt)
        center = (int(x), int(y))
        radius = int(radius)

        # 计算圆的周长与轮廓周长的比值（圆度）
        perimeter = cv2.arcLength(cnt, True)
        if perimeter == 0:
            continue
        circularity = 4 * np.pi * cv2.contourArea(cnt) / (perimeter ** 2)

        # 过滤掉非圆形的物体（圆度小于 0.8 的被忽略）
        if circularity > 0.8 and radius > 10:
            cv2.circle(frame, center, radius, (0, 255, 0), 3)  # 画圆
            cv2.circle(frame, center, 2, (0, 0, 255), 3)  # 标记圆心

            # 方向调整（让无人机朝向圆心）
            frame_center_x = frame.shape[1] // 2
            frame_center_y = frame.shape[0] // 2

            if x < frame_center_x - 50:
                print("Fly Left")
                # me.send_rc_control(-10, 0, 0, 0)
            elif x > frame_center_x + 50:
                print("Fly Right")
                # me.send_rc_control(10, 0, 0, 0)

            if y < frame_center_y - 50:
                print("Fly Up")
                # me.send_rc_control(0, 0, 10, 0)
            elif y > frame_center_y + 50:
                print("Fly Down")
                # me.send_rc_control(0, 0, -10, 0)

    cv2.imshow("Tello Camera", frame)
    cv2.imshow("Edge Detection", edges)

while True:
    vals = getKeyBoardInput()
    me.send_rc_control(vals[0], vals[1], vals[2], vals[3])

    frame = me.get_frame_read().frame
    detect_circle_by_fitting(frame)  # 处理帧
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cv2.destroyAllWindows()